#copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
#
#Licensed under the Apache License, Version 2.0 (the "License");
#you may not use this file except in compliance with the License.
#You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
#Unless required by applicable law or agreed to in writing, software
#distributed under the License is distributed on an "AS IS" BASIS,
#WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#See the License for the specific language governing permissions and
#limitations under the License.

import math
import cv2
import numpy as np
import json
import sys
import os
from PIL import Image

__all__ = ['FASTProcessTrain']


def order_box_clockwise(box):
    box = sorted(box, key=lambda x: (x[0], x[1]))
    box = np.asarray(box)
    p1 = box[0]
    d = np.linalg.norm(box[1:] - p1, axis=1)
    idx = np.argsort(d) + 1
    p2 = box[idx[0]]
    p4 = box[idx[1]]
    p3 = box[idx[2]]

    if np.cross(p1 - p2, p3 - p2) > 0:
        p2, p4 = p4, p2
    return np.asarray([p1, p2, p3, p4])

def get_rotate_matrix(alpha):
    matrix = np.array([
        [np.cos(alpha), -np.sin(alpha)],
        [np.sin(alpha), np.cos(alpha)]
    ])
    return matrix


def shrink_box(box, matrix, center, h_shrink, w_shrink):
    local_coord = matrix @ (box - center).T
    local_coord[0][local_coord[0] > 0] -= w_shrink
    local_coord[0][local_coord[0] < 0] += w_shrink
    local_coord[1][local_coord[1] > 0] -= h_shrink
    local_coord[1][local_coord[1] < 0] += h_shrink

    shrink_coord = (matrix.T @ local_coord).T + center
    return shrink_coord

def draw_rbox_map(geo_map, box, shrink_rw, shrink_rh, min_box_size=4):
    v1 = box[1] - box[0]
    v2 = box[3] - box[0]
    box_w = np.linalg.norm(v1)
    box_h = np.linalg.norm(v2)
    img_h = geo_map.shape[1]
    img_w = geo_map.shape[2]

    if np.linalg.norm(v1) >= np.linalg.norm(v2):
        theta = np.arctan2(v1[1], v1[0])
    else:
        theta = np.arctan2(v2[1], v2[0])    
        box_w, box_h = box_h, box_w
    
    max_offset = max(0, (box_h - min_box_size - 1) / 2)

    center = np.mean(box, axis=0)
    matrix = get_rotate_matrix(-theta)
    h_shrink = shrink_rh * box_h
    w_shrink = min(shrink_rw * box_w, box_h)

    if box_h < min_box_size:
        h_shrink = 0

    mask = np.zeros((img_h, img_w), dtype='uint8')
    cv2.fillPoly(mask, box[None, :], 1)
    mask = mask == 1
    xy = np.stack(np.where(mask == 1)[::-1])
    local_xy = matrix @ (xy - center.reshape(2, 1))
    geo_map[0][mask] = 0.5 * box_w + local_xy[0]
    geo_map[1][mask] = 0.5 * box_w - local_xy[0]
    geo_map[2][mask] = 0.5 * box_h + local_xy[1]
    geo_map[3][mask] = 0.5 * box_h - local_xy[1]
    geo_map[4][mask] = theta

    mask = np.zeros((img_h, img_w), dtype='uint8') 
    box1 = shrink_box(box, matrix, center, h_shrink, w_shrink)
    box1 = np.round(box1).astype('int32')
    mask = cv2.fillPoly(mask, box[None, :], 1)
    score_mask = mask == 1

    return score_mask


class FASTProcessTrain(object):
    def __init__(self, 
                 min_text_size = 4,
                 target_scale=0.5,
                 shrink_h=0.5,
                 shrink_w=0.4,
                 **kwargs):
        self.min_text_size = min_text_size
        self.target_scale = target_scale
        self.shrink_rh = shrink_h / 2
        self.shrink_rw = shrink_w / 2

    def generate_geo_map(self, img_size, polys, ignores):
        h, w = img_size
        training_mask = np.ones((h, w), dtype='uint8')
        geo_map = np.zeros((5, h, w), dtype='float32')
        score_map = np.zeros((h, w), dtype='uint8')
        polys = np.asarray(polys, dtype='int32')
        poly_areas = np.linalg.norm(polys[:, 0] - polys[:, 1], axis=1) * \
            np.linalg.norm(polys[:, 0] - polys[:, 3], axis=1)
        order = np.argsort(poly_areas)[::-1]

        for i in order:
            box = polys[i]
            if ignores[i]:
                cv2.fillPoly(training_mask, box[None, :].astype('int32'), 0)
                continue
            box = order_box_clockwise(box)
            mask = draw_rbox_map(geo_map, box, self.shrink_rw, self.shrink_rh, self.min_text_size)
            score_map[mask] = 1
        score_map = score_map.astype('float32')
        training_mask = training_mask.astype('float32')
        return score_map, geo_map, training_mask

    def cutout(self, img):
        h, w = img.shape[:2]
        ncut = np.random.randint(0, 16)
        sizes = np.random.randint(4, 32, ncut)
        ys = np.random.randint(0, h, ncut)
        xs = np.random.randint(0, w, ncut)
        vals = np.random.randint(0, 255, ncut)
        for i in range(ncut):
            img[ys[i]: ys[i] + sizes[i], xs[i]: xs[i] + sizes[i]] = vals[i]

        return img

    def __call__(self, data):
        img = data['image']
        text_polys = data['polys']
        text_tags = data['ignore_tags']
        h, w = img.shape[:2]

        if text_polys.shape[0] == 0 or np.all(text_tags):
            return None
        
        text_polys = text_polys * self.target_scale
        target_h = int(h * self.target_scale)
        target_w = int(w * self.target_scale)
        score_map, geo_map, training_mask = \
            self.generate_geo_map((target_h, target_w), text_polys, text_tags)
        geo_map[:4] /= self.target_scale

        data['score_map'] = score_map[None, :]
        data['geo_map'] = geo_map
        data['training_mask'] = training_mask[None, :]

        return data